import os
import torch
import sys
import pandas as pds
import numpy as np
from copy import deepcopy as dcopy
from collections import OrderedDict
from ..utils import ctip

## 完成
def mbox_ious(input, target):
    """计算包围盒IOU

    Args:
        input (torch.Tensor): 输入预测包围盒  #[N, 4]
        target (torch.Tensor): 输入目标包围盒 #[M, 4]

    Returns:
        torch.Tensor: 输出iou #[N, M]
    """
    M = target.shape[0]
    ix0, iy0, ix1, iy1 = torch.split(input[:,:4], 1, dim=1)
    ox0, oy0, ox1, oy1 = torch.split(target[:,:4], 1, dim=1)
    x0 = torch.max(ix0, ox0.view(-1, M))
    y0 = torch.max(iy0, oy0.view(-1, M))
    x1 = torch.min(ix1, ox1.view(-1, M))
    y1 = torch.min(iy1, oy1.view(-1, M))

    w = x1 - x0
    h = y1 - y0
    area = w * h
    area = torch.max(area, torch.zeros_like(area))

    idx = ix1 - ix0
    odx = ox1 - ox0
    idy = iy1 - iy0
    ody = oy1 - oy0

    iarea = torch.max(idx, torch.zeros_like(idx)) * torch.max(idy, torch.zeros_like(idy)) 
    oarea = torch.max(odx, torch.zeros_like(odx)) * torch.max(ody, torch.zeros_like(ody))
    sum_area = iarea + oarea.view(-1, M)
    darea = sum_area - area
    darea = torch.max(darea, torch.zeros_like(darea))
    iou = area / (darea + 1e-6)

    return iou

class mAP(object):
    """
    [A][B]: A: 表示分类正确还是错误，B: 表示预测的类别
    TP: 分类正确，分类为正样本（实际标签为正样本），P ==> P 
    TN: 分类正确，分类为负样本（实际标签为负样本），N ==> N 
    FP: 分类错误，分类为正样本（实际标签为负样本），N ==> P 误分
    FN: 分类错误，分类为负样本（实际标签为正样本），P ==> N 漏分
    """
    def __init__(self, iou_thr=[0.5], num_classes=1):
        self.num_classes = num_classes
        self.iou_thr = torch.tensor(iou_thr)
        # img_idx, x0, y0, x1, y1, conf, [class_id], gt_id, iou_0.3, iou_0.5, ... 
        self.tb = torch.zeros((0, 1+5+(num_classes>1)+1 + len(iou_thr)))

        # img_idx, x0, y0, x1, y1, [class_id] 
        self.gt = torch.zeros((0, 1+4+(num_classes>1)))
        self.table = OrderedDict([(str(k), {"TB": self.tb, "GT": self.gt}) for k in range(num_classes)])
        self.num_obj = dict((str(k), 0) for k in range(num_classes))
        self.names = []
        self.table_heads = ['img_idx', 'x0', 'y0', 'x1', 'y1', 'conf']
        if num_classes > 1:
            self.table_heads.append('class_name')
        self.table_heads.append('gt_id')
        for thr in self.iou_thr:
            self.table_heads.append('iou_'+str(round(thr.item(), 2)))
        
        self.name_dict = {}
        self.image_index = 0
        
    def __analyze(self, input:torch.Tensor, target:torch.Tensor, tag_idx):
        """计算单张图片的数值表

        Args:
            input (torch.Tensor): 单张图片、单个类别的张量[N, 2] # x0 y0 x1 y1 conf
            target (torch.Tensor): 对应的target [M, 2] # x0 y0 x1 y1
        
        Returns:
            tuple: 输出表格
        """
        img_idx = self.image_index
        np = input.shape[0]
        nt, nb = target.shape
        if np == 0 and nt != 0:
            # 漏检测 table TP, FN, FP
            gt = torch.hstack([torch.ones((nt, 1)) * img_idx, target]) 
            return self.tb, gt
        elif np != 0 and nt == 0:
            # 误检测 table TP, FN, FP
            return self.tb, self.gt 
        elif np == 0 and nt == 0:
            # 无
            return self.tb, self.gt 
        else:
            iou = mbox_ious(input, target)
            inp_iou, gtid = torch.max(iou, axis=1)
            gtid = gtid.reshape(-1, 1)
            mask = inp_iou[:, None] >= self.iou_thr

            # 增加加一行索引，用于标记样本 
            # 增加 mask和miss, 用于表示不同阈值下iou的分类情况
            table = torch.hstack([torch.ones((np, 1)) * img_idx, input, gtid + tag_idx, mask]) 
            gt = torch.hstack([torch.ones((nt, 1)) * img_idx, target])
            return table, gt 
    
    def analyze(self, input:torch.Tensor, target:torch.Tensor, name):
        """分析单张图片的情况（多个类别）

        Args:
            input (torch.Tensor): 输入包围盒 [N, 5 or 6] x0, y0, x1, y1, conf, class_id
            target (torch.Tensor): 输入GT [N, 4 or 5] x0, y0, x1, y1, class_id
            name (str or int) : 输入单张图片的标志符号
        """        
        N, S = input.shape
        try:
            assert S in [5, 6]
        except:
            ctip('red', 'bold', 'make sure your input shape is [N, 5 or 6]')
            sys.exit('abnormal exit')

        name = str(name)
        try:
            assert name not in self.name_dict.keys()
        except:
            ctip('red', 'bold', 'make sure not use this label sample[%s] before' % (name))

        if S == 5:
            # x0, y0, x1, y1, conf
            table, gt = self.__analyze(input, target, self.image_index)

            self.table['0']['TB'] = torch.vstack([self.table['0']['TB'], table])
            self.table['0']['GT'] = torch.vstack([self.table['0']['GT'], gt])
        else:
            # x0, y0, x1, y1, conf, class
            cls_idx = target[:, 4].unique()
            tid = 0
            for cid in cls_idx.long():
                input_mask = input[:,5] == cid
                cinput = input[input_mask]

                target_mask = target[:,4] == cid
                ctarget = target[target_mask]
                cid = str(cid.item())

                tid = self.table[cid]['GT'].shape[0]
                table, gt = self.__analyze(cinput, ctarget, tid)
                self.table[cid]['TB'] = torch.vstack([self.table[cid]['TB'], table])
                self.table[cid]['GT'] = torch.vstack([self.table[cid]['GT'], gt])
                self.num_obj[cid] += gt.shape[0] 
        
        #self.num_obj += target.shape[0]
        self.name_dict[str(self.image_index)] = name
        self.image_index += 1
    
    def analyzes(self, input:torch.Tensor, target:torch.Tensor):
        """批量分析图片

        Args:
            input (torch.Tensor): 第一列为图像坐标索引（从0开始）
            target (torch.Tensor): 第一列为图像坐标索引（从0开始）
        """
        apd_idx = self.image_index

        # 图像标签索引
        input_index = input[:, 0].long() + apd_idx
        target_index = target[:, 0].long() + apd_idx

        index = input_index.unique()

        try:
            assert (index == target_index.unique()).all()
        except:
            print("input", input)
            print('target', target)
            ctip('red', 'bold', '预测样本与标记样本不匹配，请核对后再尝试')
            sys.exit('异常退出')

        for idx in index:
            input_mask = (input_index == idx)
            target_mask = (target_index == idx)
            input_sub = input[input_mask]
            target_sub = target[target_mask]
            self.analyze(input_sub[:, 1:], target_sub[:, 1:], idx.item())
    
    def to_xlsx(self, path, name_dict):
        name = os.path.splitext(path)[0]
        pr_writer = pds.ExcelWriter(name+'.pr.xlsx') 
        writer = pds.ExcelWriter(path)
        table_heads = dcopy(self.table_heads)
        table_heads = ['oidx'] + table_heads        

        class_head = self.table_heads.index('class_name')
        conf_head = self.table_heads.index('conf')

        table_heads[class_head+1], table_heads[conf_head+1] = table_heads[conf_head+1], table_heads[class_head+1]

        for key in self.table.keys():
            if self.table[key]['TB'].shape[0] == 0:
                continue

            top = self.table[key]['TB'].numpy()
            top = np.concatenate([top[:, :conf_head], 
                                  top[:, class_head].reshape(-1, 1), 
                                  top[:, conf_head].reshape(-1,1), 
                                  top[:, class_head+1:]], axis=1)

            h, w = top.shape 

            mid = np.array([[''] * (w+1)]) 

            gt = self.table[key]['GT'].numpy()
            n, tw = gt.shape
            ajm = np.array([[''] * (w - tw)] * n)

            bottom = np.concatenate([gt, ajm], axis=1)

            top = np.concatenate([np.arange(h).reshape(-1, 1), top], axis=1)
            bottom = np.concatenate([np.arange(n).reshape(-1, 1), bottom], axis=1)
            data = np.concatenate([top, mid, bottom], axis=0)

            df = pds.DataFrame(data, columns=table_heads)
            
            df.to_excel(writer, sheet_name=name_dict[key], index=False)

            percision = self.pr[key]['percision']
            recall = self.pr[key]['recall']
            pr_seq = []
            pr_heads = []
            niou = self.iou_thr.shape[0]
            for i, name in enumerate(table_heads[-niou:]):
                pr_seq.append(recall[i].reshape(-1, 1))
                pr_seq.append(percision[i].reshape(-1, 1))
                pr_heads.append(name + '_r')
                pr_heads.append(name + '_p')
            data = np.concatenate(pr_seq, axis=1)
            df = pds.DataFrame(data, columns=pr_heads)
            df.to_excel(pr_writer, sheet_name=name_dict[key])


        pr_writer.save()
        writer.save()

    def __pr2ap(self, percision:list, recall:list):
        """计算ap

        Args:
            percision (list): 精度
            recall (list): 召回

        Returns:
            float: ap
        """
        ap = 0
        b = 0
        for i in range(percision.shape[0]):
            if recall[i] != b:
                mp = torch.max(percision[i:])
                ap += mp * (recall[i] - b)
                b = recall[i]

        return round(ap.item(), 4)

    def calc(self):
        pr = OrderedDict() # iou_xxx: class_xxx: [precision, recall]

        n = self.iou_thr.shape[0]
        for key in self.table.keys():
            pr[key] = {}
            table = self.table[key]['TB']
            # img_idx, x0, y0, x1, y1, conf, [class], iou_0.5, iou_0.75, ...
            conf = torch.unique(table[:, 5], sorted=True)
            conf = conf.flip(dims=[0])

            _, index = torch.sort(table[:, 5], dim=0, descending=True)
            table = table[index]
            tconf = table[:, 5]

            mask = tconf[:, None] >= conf # 阈值的mask
            mask = mask.transpose(0, 1) # 阈值(N) x 样本（M） [N, M]

            iou_ret = table[:,-n:]
            iou_ret = iou_ret.transpose(0, 1) # IOU(K) x 样本(M) [K, M]

            s = iou_ret[:, None, :] * mask[None, :, :] # IOU(K) x 阈值(N) x 样本(M) [K, N, M]

            TP = torch.sum(s, dim=2) # IOU(K) x 阈值(N) [K, N]

            recall = TP / (self.num_obj[key] + 1e-16) # IOU(K) x 阈值(N) [K, N]
            pr[key]['recall'] = torch.clip(recall, 0, 1)
            POS = torch.sum(mask, dim=1) # 输出了多少个包围盒
            percision = TP / (POS + 1e-16) # IOU(K) x 阈值(N) [K, N]
            pr[key]['percision'] = torch.clip(percision, 0, 1)
        
        ret = OrderedDict()
        for label in pr.keys():
            ret[label] = {}

            for idx, head in enumerate(self.table_heads[-n:]):
                if pr[label]['percision'].shape[1] == 0:
                    ret[label][head] = 0
                    continue
                ret[label][head] = self.__pr2ap(pr[label]['percision'][idx],
                                                pr[label]['recall'][idx])
        self.pr = pr
        self.ap = ret

        MAP = {}
        for label in self.ap:
            for key in self.ap[label].keys():
                if key not in MAP.keys():
                    MAP[key] = {"ap": 0, 'n': 0}
                MAP[key]['ap'] += self.ap[label][key]
                MAP[key]['n'] += 1
        for key in MAP.keys():
            sap = MAP[key].pop('ap')
            n = MAP[key].pop('n')

            MAP[key] = round(sap / (n + 1e-16), 4)

        return MAP

                
                
if __name__ == '__main__':
    gt_data_sets = [
        ['0001.txt',[['person',25,16,38,56],['person',129,123,41,62]]],
        ['0002.txt', [['person',123,11,43,55], ['person',38,132,59,45]]],
        ['0003.txt',[['person',16,14,35,48], ['person',123,30,49,44], ['person',99,139,47,47]]],
        ['0004.txt',[['person',53,42,40,52], ['person',154,43,31,34]]],
        ['0005.txt',[['person',59,31,44,51], ['person',48,128,34,52]]],
        ['0006.txt',[['person',36,89,52,76], ['person',62,58,44,67]]],
        ['0007.txt',[['person',28,31,55,63], ['person',58,67,50,58]]]
    ]

    det_data_sets = [
        ['0001.txt',[['person',.88,5,67,31,48],['person',.70,119,111,40,67],['person',.80,124,9,49,67]]],
        ['0002.txt', [['person',.71,64,111,64,58], ['person',.54,26,140,60,47],['person',.74,19,18,43,35]]],
        ['0003.txt',[['person',.18,109,15,77,39], ['person',.67,86,63,46,45], ['person',.38,160,62,36,53], ['person',.91,105,131,47,47], ['person',.44,18,148,40,44]]],
        ['0004.txt',[['person',.35,83,28,28,26], ['person',.78,28,68,42,67], ['person',.45,87,89,25,39], ['person',.14,10,155,60,26]]],
        ['0005.txt',[['person',.62,50,38,28,46], ['person',.44,95,11,53,28], ['person',.95,29,131,72,29], ['person',.23,29,163,72,29]]],
        ['0006.txt',[['person',.45,43,48,74,38], ['person',.84,17,155,29,35], ['person',.43,95,110,25,42]]],
        ['0007.txt',[['person',.48,16,20,101,88], ['person',.95,33,116,37,49]]]
    ]

    def _trans_gt(sets):
        name_dict = {"person": 1}
        
        gts = []
        for idx, data in enumerate(sets):
            file_name = data[0]
            for bbox in data[1]:
                class_index = name_dict[bbox[0]] 
                x0,y0,w,h = [float(x) for x in bbox[1:]]
                gts.append([idx, x0, y0, x0+w, y0+h, class_index])
            # if idx >= 1:
            #     break
        return torch.tensor(gts)        
    
    def _trans_pd(sets):
        name_dict = {"person": 1}

        pds = []
        for idx, data in enumerate(sets):
            file_name = data[0]
            for bbox in data[1]:
                class_index = name_dict[bbox[0]]
                conf = float(bbox[1])
                x0,y0,w,h = [float(x) for x in bbox[2:]]
                pds.append([idx, x0, y0, x0+w, y0+h, conf, class_index])
            # if idx >= 1:
            #     break

        return torch.tensor(pds)

    def _plot_gt_det(det, gt, img_size=(256, 256), draw=None):
        from vtorch.utils import Canvas

        if draw is None:
            draw = Canvas(img_size)
        for box in det:
            draw.rect(box[1:5].numpy(), color=(0, 255, 0), label='%f' % (box[5].numpy()))
        
        for box in gt:
            draw.rect(box[1:5].numpy(), color=(255, 0, 0), label='%d' % (box[5].numpy()))
        
        draw.show()
    
    gt = _trans_gt(gt_data_sets)
    pd = _trans_pd(det_data_sets)

    m = mAP(num_classes=2,iou_thr=[0.3, 0.5, 0.75])

    # gt1 = gt[gt[:, 0] == 6]
    # pd1 = pd[pd[:, 0] == 6]
    # print('pred', pd1)
    # print('gt', gt1)
    #a = m.push(pd1[:, 1:], gt1[:, 1:])
    from vtorch.utils import Canvas
    draw = Canvas((256, 256))

    _plot_gt_det(pd, gt, draw=draw) 
    m.analyzes(pd, gt)

    table = m.table

    for box in table['1']['TB']:
        if box[-1].bool():
            draw.rect(box[1:5].numpy(), color=(0,0,255), line_thickness=2)
        # print(box)
    
    for box in table['1']['GT']:
        if box[-1].bool():
            draw.rect(box[1:5].numpy(), color=(0, 0, 0), line_thickness=2)
    draw.show()
    
    torch.set_printoptions(sci_mode=False, precision=4)
    print(table['1']['TB'])
    print(table['1']['GT'])
    ret = m.calc()
    print(ret)

    #m.to_xlsx('test.xlsx', {"0": 'negative', '1': "person"})